## From Forum TCOD

import random
import libtcodpy as libtcod

class CityGenerator(object):
    _MAJOR_BLOCK_DIM = (60, 100)
    _MINOR_BLOCK_DIM = (15, 40)
    _BUILDING_DIM = (5, 15)

    def __init__(self, width, height):
        self.width = width
        self.height = height

    def generate(self):
        self.bsp_tree = libtcod.bsp_new_with_size(0, 0, self.width, self.height)
        # First create the major blocks
        (major_blocks, major_roads) = self._carve_blocks(self.bsp_tree,
                                          self._MAJOR_BLOCK_DIM[0],
                                          self._MAJOR_BLOCK_DIM[1], 8)
        # Second create the minor blocks
        minor_blocks = []
        minor_roads = []
        for major_node in major_blocks:
            (blocks, roads) = self._carve_blocks(major_node,
                                                 self._MINOR_BLOCK_DIM[0],
                                                 self._MINOR_BLOCK_DIM[1], 4)
            minor_blocks.extend(blocks)
            minor_roads.extend(roads)
        # Finally create the building sites
        buildings = []
        alleys = []
        for minor_node in minor_blocks:
            (blocks, roads) = self._carve_blocks(minor_node,
                                                 self._BUILDING_DIM[0],
                                                 self._BUILDING_DIM[1], 1)
            buildings.extend(blocks)
            alleys.extend(roads)

    def _carve_blocks(self, root_node, min_leaf_dim, max_leaf_dim, road_width):
        '''Carve the node into number blocks.

        Returns the (blocks, roads).
        '''
        orientation = random.randint(0, 1)
        previous_splits = True
        roads = []
        while previous_splits:
            previous_splits = 0
            leaves = self._get_leaves(root_node)
            # Split the blocks
            for node in leaves:
                # Check the max_leaf_dim
                if max(node.w, node.h) < random.randrange(min_leaf_dim,
                                                          max_leaf_dim):

                    continue
                # Horizontal split
                if not orientation:
                    leftmost = int(node.x + node.w*0.4)
                    rightmost = int(node.x + node.w*0.6)
                    if leftmost == rightmost:
                        position = leftmost
                    else:
                        position = random.randrange(leftmost, rightmost)
                    smallest_dim = min(node.x + node.w - position,
                                       position - node.x)
                # Vertical split
                else:
                    topmost = int(node.y + node.h*0.4)
                    bottommost = int(node.y + node.h*0.6)
                    if topmost == bottommost:
                        position = topmost
                    else:
                        position = random.randrange(topmost, bottommost)
                    smallest_dim = min(node.y + node.h - position,
                                       position - node.y)
                if smallest_dim < min_leaf_dim:
                    continue
                libtcod.bsp_split_once(node, orientation, position)
                previous_splits += 1
                # Carve out the splitting road
                splitting_road = self._carve_road(node, orientation, road_width)
                roads.append(splitting_road)
            # Flip the orientation for the next split
            orientation = not orientation
        return (self._get_leaves(root_node), roads)

    def _carve_road(self, node, orientation, road_width):
        '''Carve out the splitting roads.
        '''
        if (road_width % 2) == 0:
            left_portion = road_width / 2
            right_portion = road_width / 2
        else:
            randint = random.randint(0, 1)
            left_portion = road_width / 2 + randint
            right_portion = road_width / 2 + abs(1 - randint)
        left_child = libtcod.bsp_left(node)
        right_child = libtcod.bsp_right(node)
        if not orientation:
            top_left = (node.position - left_portion, node.y)
            bottom_right = (node.position + right_portion, node.y + node.h)
            libtcod.bsp_resize(left_child, left_child.x, left_child.y,
                               left_child.w - left_portion, left_child.h)
            libtcod.bsp_resize(right_child, right_child.x + right_portion,
                               right_child.y,
                               right_child.w - right_portion, right_child.h)
        else:
            top_left = (node.x, node.position - left_portion)
            bottom_right = (node.x + node.w, node.position + right_portion)
            libtcod.bsp_resize(left_child, left_child.x, left_child.y,
                               left_child.w, left_child.h - left_portion)
            libtcod.bsp_resize(right_child, right_child.x,
                               right_child.y + right_portion,
                               right_child.w, right_child.h - right_portion)
        return (top_left, bottom_right)

    def _get_leaves(self, root_node=None):
        if root_node is None:
            root_node = self.bsp_tree
        leaves = []
        def leaf_callback(node, userData):
            if libtcod.bsp_is_leaf(node):
                leaves.append(node)
            return True
        libtcod.bsp_traverse_pre_order(root_node, leaf_callback)
        return leaves


def demo(city_width, city_height):
    city_generator = CityGenerator(city_width, city_height)
    libtcod.console_init_root(200, 200, 'CityGenerator Demo', False)
    running = True
    regen = True
    colors = {}
    offset = [0, 0]
    while running:
        if regen:
            city_generator.generate()
            regen = False
            colors = {}
        colors = _render(city_generator.bsp_tree, offset, colors)
        key = libtcod.console_wait_for_keypress(True)
        if key.vk == libtcod.KEY_ESCAPE:
            running = False
        elif key.vk == libtcod.KEY_SPACE:
            regen = True
        elif key.vk == libtcod.KEY_LEFT:
            offset[0] += 10
        elif key.vk == libtcod.KEY_RIGHT:
            offset[0] -= 10
        elif key.vk == libtcod.KEY_UP:
            offset[1] += 10
        elif key.vk == libtcod.KEY_DOWN:
            offset[1] -= 10
        elif key.c == ord('s'):
            libtcod.sys_save_screenshot()

def _render(bsp_tree, offset=(0, 0), colors={}):
    libtcod.console_clear(None)
    class GeneratorRender(object):
        def __init__(self, colors={}):
            self.node_colors = colors

        def next_color(self, node):
            node_id = (node.x, node.y, node.w, node.h)
            if node_id in self.node_colors:
                cached_color = self.node_colors[node_id]
            else:
                cached_color = libtcod.color_lerp(libtcod.white,
                                                  libtcod.dark_grey,
                                                  random.random())
                self.node_colors[node_id] = cached_color
            return cached_color

        def __call__(self, node, userData):
            if libtcod.bsp_is_leaf(node):
                node_color = self.next_color(node)
                for x in range(node.x, node.x + node.w):
                    for y in range(node.y, node.y + node.h):
                        libtcod.console_set_back(None, x + offset[0],
                                                 y + offset[1], node_color)
            return True
    render_callback = GeneratorRender(colors)
    libtcod.bsp_traverse_in_order(bsp_tree, render_callback)
    libtcod.console_flush()
    return render_callback.node_colors


if __name__ == "__main__":
    demo(200, 200)
